#-*- codig:utf-8 -*-
import akshare as ak
from datetime import datetime, timedelta
import time
import numpy as np
import math
import pandas as pd
from get_data import get_price, get_index_price, get_security_info, get_trade_days

# https://www.joinquant.com/help/api/help#api:API%E6%96%87%E6%A1%A3
# https://www.joinquant.com/help/api/help#JQData:JQData

# 请替换为您的聚宽账号和密码
# jq.auth('18589010884','131421')

# http://fund.eastmoney.com/ETFN_jzzzl.html
stock_pool = [
    '159915.XSHE', # 易方达创业板ETF
    '510300.XSHG', # 华泰柏瑞沪深300ETF
    '510500.XSHG', # 南方中证500ETF
]

# 动量轮动参数
stock_num = 1           # 买入评分最高的前 stock_num 只股票
momentum_day = 29       # 最新动量参考最近 momentum_day 的

ref_stock = '000300.XSHG' #用 ref_stock 做择时计算的基础数据
N = 18 # 计算最新斜率 slope，拟合度 r2 参考最近 N 天
M = 120 # 计算最新标准分 zscore，rsrs_score 参考最近 M 天（降低到120天）
score_threshold = 0.7 # rsrs 标准分指标阈值
# ma 择时参数
mean_day = 20 # 计算结束 ma 收盘价，参考最近 mean_day
mean_diff_day = 3 # 计算初始 ma 收盘价，参考(mean_day + mean_diff_day)天前，窗口为 mean_diff_day 的一段时间


day = 1

# 根据股票名，获取股票 code
def get_stock_code(stock_name):
    # 简化版本，直接返回股票代码
    # 在实际使用中，可以通过其他方式获取股票代码
    return stock_name

# # 根据市值，获取股票池
# def market_cap(): 
#     wholeA = jq.get_fundamentals(jq.query(
#         jq.valuation.code).filter(
#             jq.valuation.market_cap > 2000
#         ))
#     wholeAList = list(wholeA['code'])
#     return wholeAList

# 1-1 选股模块-动量因子轮动 
# 基于股票年化收益和判定系数打分,并按照分数从大到小排名
def get_rank(stock_pool, end_date=None):
    score_list = []
    
    # 计算基准日期
    if end_date is None:
        current_dt = time.strftime("%Y-%m-%d", time.localtime())
        current_dt = datetime.strptime(current_dt, '%Y-%m-%d')
        end_date = current_dt - timedelta(days=day)
    
    # 确保结束日期不超过权限范围
    if end_date > datetime(2025, 6, 17):
        end_date = datetime(2025, 6, 17)
    
    print(f"  动量计算基准日期: {end_date.strftime('%Y-%m-%d')}")
    
    for stock in stock_pool:
        try:
            # 获取股票价格数据
            data = get_price(stock, end_date=end_date, count=momentum_day, 
                              frequency='daily', fields=['close'])
            
            
            if len(data) < momentum_day:
                print(f"  ⚠️  {stock} 数据不足: 需要{momentum_day}天，实际{len(data)}天")
                score_list.append(0)  # 数据不足时给0分
                continue
            
            # 计算对数收益率
            if isinstance(data, pd.DataFrame):
                close_values = data['close'].values if 'close' in data.columns else data.iloc[:, 0].values
            else:
                close_values = np.asarray(data)
            
            y = np.log(close_values)
            x = np.arange(len(y))
            
            # 线性回归拟合
            slope, intercept = np.polyfit(x, y, 1)
            
            # 计算年化收益率
            annualized_returns = math.pow(math.exp(slope), 250) - 1
            
            # 计算拟合度R²
            y_pred = slope * x + intercept
            ss_res = np.sum((y - y_pred) ** 2)
            ss_tot = np.sum((y - np.mean(y)) ** 2)
            r_squared = 1 - (ss_res / ss_tot) if ss_tot != 0 else 0
            
            # 综合评分
            score = annualized_returns * r_squared
            score_list.append(score)
            
            # 获取股票名称用于显示
            security_info = get_security_info(stock)
            stock_name = security_info.display_name
            print(f"  {stock_name}({stock}): 年化收益={annualized_returns:.4f}, R²={r_squared:.4f}, 评分={score:.4f}")
            
        except Exception as e:
            print(f"  ❌ {stock} 计算失败: {str(e)}")
            score_list.append(0)
    
    # 创建股票-分数字典并排序
    stock_dict = dict(zip(stock_pool, score_list))
    sort_list = sorted(stock_dict.items(), key=lambda item: item[1], reverse=True)
    
    # 返回排名前stock_num的股票
    rank_stock = [sort_list[i][0] for i in range(min(stock_num, len(sort_list)))]
    return rank_stock

# 2-1 择时模块-计算线性回归统计值
# 对输入的自变量每日最低价 x(series) 和因变量每日最高价 y(series) 建立 OLS 回归模型,返回元组(截距,斜率,拟合度)
# R2 统计学线性回归决定系数，也叫判定系数，拟合优度。
# R2 范围 0 ~ 1，拟合优度越大，自变量对因变量的解释程度越高，越接近 1 越好。
# 公式说明： https://blog.csdn.net/snowdroptulip/article/details/79022532
#           https://www.cnblogs.com/aviator999/p/10049646.html
def get_ols(x, y):
    slope, intercept = np.polyfit(x, y, 1)
    r2 = 1 - (sum((y - (slope * x + intercept))**2) / ((len(y) - 1) * np.var(y, ddof=1)))
    return (intercept, slope, r2)

# 2-2 择时模块-设定初始斜率序列
# 通过前 M 日最高最低价的线性回归计算初始的斜率,返回斜率的列表
def initial_slope_series(end_date=None):
    
        # 方法1：使用 datetime.now() 直接获取当前时间
    current_dt = datetime.now()
    start_date = current_dt - timedelta(days=M)
    
    # 格式化日期为字符串用于API调用
    start_date = start_date.strftime('%Y%m%d')
    end_date = current_dt.strftime('%Y%m%d')
   
        
    # 使用日期范围而不是count参数
    # start_date = '20240610'
    #  '000300.XSHG' #用 ref_stock 做择时计算的基础数据 转成 sh000300
    # 将 ref_stock 赋值为 '000300.XSHG'，并将 day、N、M 设置为常用参数
    ref_stock = 'sh000300'
    print(start_date,end_date)
    data = get_index_price(ref_stock, start_date=start_date, end_date=end_date)
    print(data)
    # data = data[['high', 'low']]
    # 确保数据量足够进行计算
    if len(data) < N:
        raise Exception("数据不足，无法进行计算")
    elif len(data) < N + M:
        # 如果数据不够M+N天，就使用全部可用数据
        required_data = data
    else:
        # 如果数据足够，取最近的N+M条数据
        required_data = data.tail(N + M)
        
    return [get_ols(required_data.low.iloc[i:i+N], required_data.high.iloc[i:i+N])[1] for i in range(len(required_data) - N)]

# 2-3 择时模块-计算标准分
# 通过斜率列表计算并返回截至回测结束日的最新标准分
def get_zscore(slope_series):
    mean = np.mean(slope_series)
    std = np.std(slope_series)
    return (slope_series[-1] - mean) / std

# 2-4 择时模块-计算综合信号
# 1.获得 rsrs 与 MA 信号,rsrs 信号算法参考优化说明，MA 信号为一段时间两个端点的 MA 数值比较大小
# 2.信号同时为 True 时返回买入信号，同为 False 时返回卖出信号，其余情况返回持仓不变信号
# 解释：
#       MA 信号：MA 指标是英文(Moving average)的简写，叫移动平均线指标。
#       RSRS 择时信号：
#               https://www.joinquant.com/view/community/detail/32b60d05f16c7d719d7fb836687504d6?type=1
def get_timing_signal(stock, end_date=None):
    global slope_series
    current_dt = datetime.now()
    start_date = current_dt - timedelta(days=M)
    
    # 格式化日期为字符串用于API调用
    start_date = start_date.strftime('%Y%m%d')
    end_date = current_dt.strftime('%Y%m%d')

    
    try:
        # 获取截至当前日期的收盘价数据计算MA
        close_data = get_price(ref_stock, start_date=start_date, end_date=end_date, 
                                frequency='daily', fields=['close'])
        
        # 确保数据量足够
        if len(close_data) < mean_day + mean_diff_day:
            print(f"⚠️  MA数据不足: 需要{mean_day + mean_diff_day}天，实际{len(close_data)}天")
            return "KEEP"
        
        # 计算移动平均线（基于当前日期的数据）
        today_MA = close_data.close[-mean_day:].mean() 
        before_MA = close_data.close[-(mean_day + mean_diff_day):-mean_diff_day].mean()
        
        # 获取截至当前日期的高低价数据计算RSRS
        high_low_data = get_price(ref_stock, start_date=start_date, end_date=calculation_date, 
                                   frequency='daily', fields=['high', 'low'])
        
        if len(high_low_data) < N:
            print(f"⚠️  RSRS数据不足: 需要{N}天，实际{len(high_low_data)}天")
            return "KEEP"
            
        # 计算当前日期的RSRS指标
        recent_high_low = high_low_data.tail(N)
        intercept, slope, r2 = get_ols(recent_high_low.low, recent_high_low.high)
        
        # 更新斜率序列（添加当前日期的斜率）
        slope_series.append(slope)
        
        # 动态调整M值，确保有足够的历史数据
        available_data = len(slope_series)
        effective_M = min(M, available_data) if available_data >= 30 else available_data
        
        if effective_M < 30:  # 至少需要30个数据点
            print(f"⚠️  斜率历史数据不足: 需要至少30个，实际{available_data}个")
            return "KEEP"
            
        # 计算RSRS分数（使用动态调整的M值）
        rsrs_score = get_zscore(slope_series[-effective_M:]) * r2
        
        # 输出调试信息
        print(f"  MA信号: 当前={today_MA:.4f}, 之前={before_MA:.4f}, 趋势={'上升' if today_MA > before_MA else '下降'}")
        print(f"  RSRS信号: 分数={rsrs_score:.4f}, 阈值=±{score_threshold}, R²={r2:.4f}, 使用{effective_M}天数据")
        
        # 综合判断信号
        ma_signal = today_MA > before_MA
        rsrs_buy_signal = rsrs_score > score_threshold
        rsrs_sell_signal = rsrs_score < -score_threshold
        
        if rsrs_buy_signal and ma_signal:
            return "BUY"
        elif rsrs_sell_signal and not ma_signal:
            return "SELL"
        else:
            return "KEEP"
            
    except Exception as e:
        print(f"❌ 择时信号计算失败: {str(e)}")
        return "KEEP"

slope_series = initial_slope_series()[:-1] # 除去回测第一天的 slope ，避免运行时重复加入


def run_single_day():
    """运行单日策略"""
    global slope_series
    
    print("=== 今日策略运行 ===")
    
    try:
        # 初始化斜率序列
        slope_series = initial_slope_series()[:-1]
        print(f"✅ 斜率序列初始化完成，长度: {len(slope_series)}")
        
        # 获取选股结果（使用最新数据）
        check_out_list = get_rank(stock_pool)
        
        print("\n📊 今日股票排名:")
        for i, stock_code in enumerate(check_out_list):
            print(stock_code)
            security_info = get_security_info(stock_code)
            stock_name = security_info.display_name
            print(f'  {i+1}. {stock_name}({stock_code})')
        
        # 获取择时信号
        timing_signal = get_timing_signal(ref_stock)
        print(f'\n🎯 择时信号: {timing_signal}')
        
        # 输出操作建议
        if timing_signal == "BUY" and check_out_list:
            top_stock = check_out_list[0]
            security_info = get_security_info(top_stock)
            print(f'💡 操作建议: 买入 {security_info.display_name}({top_stock})')
        elif timing_signal == "SELL":
            print('💡 操作建议: 卖出当前持仓')
        else:
            print('💡 操作建议: 保持当前仓位')
            
    except Exception as e:
        print(f'❌ 策略运行失败: {str(e)}')
    
    print('=' * 50)

def quick_test():
    """快速测试最近几天的策略表现"""
    global slope_series
    
    print("=== 快速回测（最近5个交易日）===")
    
    try:
        # 获取最近的交易日
        end_date = datetime(2025, 6, 17)
        trade_days = get_trade_days(end_date=end_date, count=5)
        
        # 初始化斜率序列（使用第一个交易日的数据）
        init_date = trade_days[0]
        if isinstance(init_date, datetime):
            init_datetime = init_date
        else:
            init_datetime = datetime.combine(init_date, datetime.min.time())
        
        slope_series = initial_slope_series(init_datetime)[:-1]
        
        for i, trade_date in enumerate(trade_days):
            # 确保trade_date是datetime类型
            if isinstance(trade_date, datetime):
                trade_datetime = trade_date
            else:
                trade_datetime = datetime.combine(trade_date, datetime.min.time())
            
            print(f"\n--- {trade_date.strftime('%Y-%m-%d')} ---")
            
            # 选股（使用当前交易日的数据）
            check_out_list = get_rank(stock_pool, trade_datetime)
            top_stock = check_out_list[0] if check_out_list else None
            
            if top_stock:
                security_info = get_security_info(top_stock)
                print(f"推荐股票: {security_info.display_name}")
            
            # 择时（使用当前交易日的数据）
            timing_signal = get_timing_signal(ref_stock, trade_datetime)
            print(f"择时信号: {timing_signal}")
            
    except Exception as e:
        print(f"❌ 快速测试失败: {str(e)}")
    
    print('=' * 50)

if __name__ == "__main__":
    # 可以选择运行单日策略或回测
    import sys
    run_single_day()